{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2020-05-04T07:00:02.811268Z",
          "start_time": "2020-05-04T07:00:02.808081Z"
        },
        "cellView": "form",
        "id": "zh8WOre17mel"
      },
      "outputs": [],
      "source": [
        "#@title ##### License\n",
        "# Copyright 2018 The GraphNets Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#    http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# ============================================================================"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rPPRpGpz1Hce"
      },
      "source": [
        "![](https://github.com/deepmind/graph_nets/raw/master/images/graph-nets-deepmind-shortest-path.gif)\n",
        "\n",
        "# Tutorial of the Graph Nets library\n",
        "\n",
        "The purpose of this tutorial is to get practical experience using the Graph Nets library via examples of:\n",
        "1. Building graph data structures `graph_nets.graphs.GraphsTuple` using `graph_nets.utils_np` .\n",
        "2. Operating with graph data structures in the tensorflow graph using `graph_nets.utils_tf` .\n",
        "3. Feeding graphs to Graph Nets tensorflow modules in `graph_nets.modules` .\n",
        "4. Building custom Graph Nets modules using the graph net building blocks provided in `graph_nets.blocks` .\n",
        "\n",
        "For more information about graph networks, see our arXiv paper: [Relational inductive biases, deep learning, and graph networks](https://arxiv.org/abs/1806.01261)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2020-05-04T07:00:03.697364Z",
          "start_time": "2020-05-04T07:00:03.691059Z"
        },
        "cellView": "form",
        "id": "bq24M0Ig7r8N"
      },
      "outputs": [],
      "source": [
        "#@title ### Install the Graph Nets library on this Colaboratory runtime  { form-width: \"60%\", run: \"auto\"}\n",
        "#@markdown \u003cbr\u003e1. Connect to a local or hosted Colaboratory runtime by clicking the **Connect** button at the top-right.\u003cbr\u003e2. Choose \"Yes\" below to install the Graph Nets library on the runtime machine with the correct dependencies. Note, this works both with local and hosted Colaboratory runtimes.\n",
        "\n",
        "install_graph_nets_library = \"No\"  #@param [\"Yes\", \"No\"]\n",
        "\n",
        "if install_graph_nets_library.lower() == \"yes\":\n",
        "  print(\"Installing Graph Nets library and dependencies:\")\n",
        "  print(\"Output message from command:\\n\")\n",
        "  !pip install \"graph_nets\u003e=1.1\" \"dm-sonnet\u003e=2.0.0b0\" \"tensorflow_probability\"\n",
        "else:\n",
        "  print(\"Skipping installation of Graph Nets library\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1S-8wRohRDE2"
      },
      "source": [
        "### Install dependencies locally\n",
        "\n",
        "If you are running this notebook locally (i.e., not through Colaboratory), you will also need to install a few more dependencies. Run the following on the command line to install the graph networks library, as well as a few other dependencies:\n",
        "\n",
        "```\n",
        "pip install \"graph_nets\u003e=1.1\" \"tensorflow\u003e=2.1.0-rc1\" \"dm-sonnet\u003e=2.0.0b0\" tensorflow_probability\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2020-05-04T07:00:12.053476Z",
          "start_time": "2020-05-04T07:00:07.223206Z"
        },
        "cellView": "form",
        "id": "ww1nT7Ddydrk",
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "#@title #### (Imports)\n",
        "%tensorflow_version 2.x  # For Google Colab only.\n",
        "\n",
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "from graph_nets import blocks\n",
        "\n",
        "from graph_nets import graphs\n",
        "from graph_nets import modules\n",
        "from graph_nets import utils_np\n",
        "from graph_nets import utils_tf\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import networkx as nx\n",
        "import numpy as np\n",
        "import sonnet as snt\n",
        "import tensorflow as tf\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "toc",
        "id": "PbT9ciFt4vGk"
      },
      "source": [
        "\u003e[Tutorial of the Graph Nets library](#scrollTo=rPPRpGpz1Hce)\n",
        "\n",
        "\u003e\u003e\u003e[Install dependencies locally](#scrollTo=1S-8wRohRDE2)\n",
        "\n",
        "\u003e\u003e[The graphs.GraphsTuple class](#scrollTo=Tc355RxEak3h)\n",
        "\n",
        "\u003e[Creating graphs](#scrollTo=Rg3HRdEAyYi0)\n",
        "\n",
        "\u003e\u003e[What's contained in a graph?](#scrollTo=VC023vAp0pAp)\n",
        "\n",
        "\u003e\u003e[How to represent graphs as a graphs.GraphsTuple](#scrollTo=8sOteecpHj0Z)\n",
        "\n",
        "\u003e\u003e\u003e[Visualize the graphs using networkx](#scrollTo=saTlW6FUsQou)\n",
        "\n",
        "\u003e\u003e\u003e[Print the GraphsTuple fields](#scrollTo=o-HixEQ7sU6q)\n",
        "\n",
        "\u003e\u003e\u003e[Back to data dicts](#scrollTo=3TaUiWVh_SX9)\n",
        "\n",
        "\u003e\u003e[Ways to represent different data sources with a graph](#scrollTo=JtUYYeERq2x2)\n",
        "\n",
        "\u003e\u003e\u003e[Graph with no features](#scrollTo=nz7RDJL1r8Mh)\n",
        "\n",
        "\u003e\u003e\u003e[Set (ie. graph without edges)](#scrollTo=efQtkXoFr5su)\n",
        "\n",
        "\u003e\u003e\u003e[Creating a GraphsTuple from a networkx graph](#scrollTo=oPrZFvcRr1yU)\n",
        "\n",
        "\u003e\u003e[Working with tensor GraphsTuple's](#scrollTo=cSUY6pbYXE1i)\n",
        "\n",
        "\u003e\u003e[Creating a constant tensor GraphsTuple from data dicts](#scrollTo=_kygzbsSXsHJ)\n",
        "\n",
        "\u003e\u003e[Slicing graphs from within a batch](#scrollTo=acSavUGPXsZI)\n",
        "\n",
        "\u003e\u003e[Concatenating multiple GraphsTuple instances](#scrollTo=MLD2KEdEYldQ)\n",
        "\n",
        "\u003e[Graph Net modules](#scrollTo=VUzXQDdRaIKt)\n",
        "\n",
        "\u003e\u003e[Creating a modules.GraphNetwork](#scrollTo=0rOU3rjRbK9O)\n",
        "\n",
        "\u003e\u003e[Feeding a GraphsTuple to a Graph Net](#scrollTo=VsSIrvMEbV1W)\n",
        "\n",
        "\u003e\u003e[Connecting a GraphNetwork recurrently](#scrollTo=HGK5xGBob3lb)\n",
        "\n",
        "\u003e\u003e[Other canonical Graph Net modules](#scrollTo=6m6aOhOtAz3T)\n",
        "\n",
        "\u003e[Graph Net building blocks](#scrollTo=qfhlV2qpah5O)\n",
        "\n",
        "\u003e\u003e[Broadcast operations](#scrollTo=Ta0YenqKaq7E)\n",
        "\n",
        "\u003e\u003e[Aggregators](#scrollTo=uE8bBiLbaoCX)\n",
        "\n",
        "\u003e\u003e[blocks.EdgeBlock](#scrollTo=lE0e3bMna96z)\n",
        "\n",
        "\u003e\u003e[blocks.NodeBlock](#scrollTo=FsqdEWjLa0kF)\n",
        "\n",
        "\u003e\u003e[blocks.GlobalBlock](#scrollTo=JlNdE_rbbBOj)\n",
        "\n",
        "\u003e\u003e[Block compositionality](#scrollTo=ZC_Y6dEZqNCx)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tc355RxEak3h"
      },
      "source": [
        "## The [`graphs.GraphsTuple`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-graphsgraphstuple) class\n",
        "\n",
        "The Graph Nets library contains models which operate on graph-structured data, so the first thing to understand is how graph-structured data is represented in the code. \n",
        "\n",
        "The `graph_nets.graphs.GraphsTuple` class, defined in [`graph_nets/graphs.py`](https://github.com/deepmind/graph_nets/blob/master/graph_nets/graphs.py), represents a **batches of one or more graphs**. All graph network modules take instances of `GraphsTuple` as input, and return instances of `GraphsTuple` as output. The graphs are *directed* (one-way edges), *attributed* (node-, edge-, and graph-level features are allowed), *multigraphs* (multiple edges can connect any two nodes, and self-edges are allowed). See Box 3, page 11 in [our companion arXiv paper](https://arxiv.org/pdf/1806.01261.pdf) for details.\n",
        "\n",
        "A `GraphsTuple` has attributes:\n",
        "- `n_node` (shape=[num_graphs]): Number of nodes in each graph in the batch.\n",
        "- `n_edge` (shape=[num_graphs]): Number of edges in each graph in the batch.\n",
        "- `globals` (shape=[num_graphs] + global_feature_dimensions): Global features for each graph in the batch.\n",
        "- `nodes` (shape=[total_num_nodes] + node_feature_dimensions): Node features for each node in the batch of graphs.\n",
        "- `edges` (shape=[total_num_edges] + edge_feature_dimensions): Edge features for each edge in the batch of graphs.\n",
        "- `senders`(shape=[total_num_edges]): Indices of the nodes in `nodes`, which indicate the source node of each directed edge in `edges`.\n",
        "- `receivers` (shape=[total_num_edges]): Indices of the nodes in `nodes`, which indicate the destination node of each directed edge in `edges`.\n",
        "\n",
        "The nodes and edges from the different graphs in the batch are concatenated along the first axis of the `nodes` and `edges` fields, and can be partitioned using the `n_node` and `n_edge` fields respectively.  Note, all but the \"`n_*`\" fields are optional (see examples below).\n",
        "\n",
        "The attributes of a `GraphsTuple` instance are typically either Numpy arrays or TensorFlow tensors. The library contains utilities for manipulating graphs with each of these types of attributes, respectively:\n",
        "  * `utils_np` (for Numpy arrays)\n",
        "  * `utils_tf` (for TensorFlow tensors)\n",
        "\n",
        "An important method of the `GraphsTuple` class is [`GraphsTuple.replace`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#graphsgraphstuplereplacekwargs): Similarly to `collections.namedtuple._replace` (in fact, `GraphsTuple` is sub-class of `collections.namedtuple`), this method creates a copy of the `GraphsTuple`, with references to all of the original attributes, by replacing some of them by the values provided as keyword arguments."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rg3HRdEAyYi0"
      },
      "source": [
        "# Creating graphs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VC023vAp0pAp"
      },
      "source": [
        "## What's contained in a graph?\n",
        "\n",
        "Each graph will have a global feature, several nodes, and several edges. The graphs can have different numbers of nodes and edges, but the lengths of the global, node, and edge attribute vectors must be the same across graphs.\n",
        "In order to create a `graphs.GraphsTuple` instance, we can define a `list` whose elements are `dict`s, with the following keys, that contain each graph's data:\n",
        "- \"globals\": Each graph has a single `float`-valued feature vector.\n",
        "- \"nodes\": Each graph has a set of nodes with `float`-valued feature vectors.\n",
        "- \"edges\": Each graph has a set of edges with `float`-valued feature vectors.\n",
        "- \"senders\": Each edge connects a sender node, represented by an `int`-valued node index, to a receiver node.\n",
        "- \"receivers\": Each edge connects a sender node to a receiver node, represented by an `int`-valued node index.\n",
        "\n",
        "Try running the cell below to create some dummy graph data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2020-05-04T06:59:11.766943Z",
          "start_time": "2020-05-04T06:59:01.119Z"
        },
        "cellView": "both",
        "id": "jW6MIo92HgOR"
      },
      "outputs": [],
      "source": [
        "# Global features for graph 0.\n",
        "globals_0 = [1., 2., 3.]\n",
        "\n",
        "# Node features for graph 0.\n",
        "nodes_0 = [[10., 20., 30.],  # Node 0\n",
        "           [11., 21., 31.],  # Node 1\n",
        "           [12., 22., 32.],  # Node 2\n",
        "           [13., 23., 33.],  # Node 3\n",
        "           [14., 24., 34.]]  # Node 4\n",
        "\n",
        "# Edge features for graph 0.\n",
        "edges_0 = [[100., 200.],  # Edge 0\n",
        "           [101., 201.],  # Edge 1\n",
        "           [102., 202.],  # Edge 2\n",
        "           [103., 203.],  # Edge 3\n",
        "           [104., 204.],  # Edge 4\n",
        "           [105., 205.]]  # Edge 5\n",
        "\n",
        "# The sender and receiver nodes associated with each edge for graph 0.\n",
        "senders_0 = [0,  # Index of the sender node for edge 0\n",
        "             1,  # Index of the sender node for edge 1\n",
        "             1,  # Index of the sender node for edge 2\n",
        "             2,  # Index of the sender node for edge 3\n",
        "             2,  # Index of the sender node for edge 4\n",
        "             3]  # Index of the sender node for edge 5\n",
        "receivers_0 = [1,  # Index of the receiver node for edge 0\n",
        "               2,  # Index of the receiver node for edge 1\n",
        "               3,  # Index of the receiver node for edge 2\n",
        "               0,  # Index of the receiver node for edge 3\n",
        "               3,  # Index of the receiver node for edge 4\n",
        "               4]  # Index of the receiver node for edge 5\n",
        "\n",
        "# Global features for graph 1.\n",
        "globals_1 = [1001., 1002., 1003.]\n",
        "\n",
        "# Node features for graph 1.\n",
        "nodes_1 = [[1010., 1020., 1030.],  # Node 0\n",
        "           [1011., 1021., 1031.]]  # Node 1\n",
        "\n",
        "# Edge features for graph 1.\n",
        "edges_1 = [[1100., 1200.],  # Edge 0\n",
        "           [1101., 1201.],  # Edge 1\n",
        "           [1102., 1202.],  # Edge 2\n",
        "           [1103., 1203.]]  # Edge 3\n",
        "\n",
        "# The sender and receiver nodes associated with each edge for graph 1.\n",
        "senders_1 = [0,  # Index of the sender node for edge 0\n",
        "             0,  # Index of the sender node for edge 1\n",
        "             1,  # Index of the sender node for edge 2\n",
        "             1]  # Index of the sender node for edge 3\n",
        "receivers_1 = [0,  # Index of the receiver node for edge 0\n",
        "               1,  # Index of the receiver node for edge 1\n",
        "               0,  # Index of the receiver node for edge 2\n",
        "               0]  # Index of the receiver node for edge 3\n",
        "\n",
        "data_dict_0 = {\n",
        "    \"globals\": globals_0,\n",
        "    \"nodes\": nodes_0,\n",
        "    \"edges\": edges_0,\n",
        "    \"senders\": senders_0,\n",
        "    \"receivers\": receivers_0\n",
        "}\n",
        "\n",
        "data_dict_1 = {\n",
        "    \"globals\": globals_1,\n",
        "    \"nodes\": nodes_1,\n",
        "    \"edges\": edges_1,\n",
        "    \"senders\": senders_1,\n",
        "    \"receivers\": receivers_1\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8sOteecpHj0Z"
      },
      "source": [
        "## How to represent graphs as a [`graphs.GraphsTuple`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#utils_npgraphs_tuple_to_data_dictsgraph)\n",
        "\n",
        "The `utils_np` module contains a functions named [`utils_np.data_dicts_to_graphs_tuple`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#utils_npgraphs_tuple_to_data_dictsgraph), which takes a `list` of `dict`s with the keys specified above, and returns a `GraphsTuple` that represents the sequence of graphs.\n",
        "\n",
        "The `data_dicts_to_graphs_tuple` function does three things:\n",
        "1. It concatenates the data from the multiple graphs together along their innermost axes (i.e. the batch dimension). This allows the graph net to process the node and edge attributes by a shared function in parallel.\n",
        "2. It counts up the numbers of nodes and edges per graph, and stores them in the fields \"n_node\" and \"n_edge\", respectively, where their lengths equal the number of graphs. This is used for keeping track of which nodes and edges belong to which graph, so they can be split up later, and so the graph can broadcast a graph's global attributes across its nodes and edges.\n",
        "3. It adds an integer offset to the sender and receiver indices, which corresponds to the number of nodes in the preceding graphs. This allows the indices to correspond to the nodes and edges of their corresponding graph, after the node and edge attributes have been concatenated.\n",
        "\n",
        "Try running the cell below to put the graph dictionaries into a `GraphsTuple` using `utils_np.data_dicts_to_graphs_tuple`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2020-05-04T06:59:11.768906Z",
          "start_time": "2020-05-04T06:59:04.342Z"
        },
        "id": "cpUOMJo30oHa"
      },
      "outputs": [],
      "source": [
        "data_dict_list = [data_dict_0, data_dict_1]\n",
        "graphs_tuple = utils_np.data_dicts_to_graphs_tuple(data_dict_list)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "saTlW6FUsQou"
      },
      "source": [
        "### Visualize the graphs using `networkx`\n",
        "\n",
        "A `GraphsTuple` can be converted into a `list` of `networkx` graph objects for easy visualization.\n",
        "\n",
        "Try running the cell below to visualize the graphs we've just defined."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2020-05-04T06:59:13.706735Z",
          "start_time": "2020-05-04T06:59:13.685161Z"
        },
        "cellView": "both",
        "id": "B-vAc0An9Dvu"
      },
      "outputs": [],
      "source": [
        "graphs_nx = utils_np.graphs_tuple_to_networkxs(graphs_tuple)\n",
        "_, axs = plt.subplots(ncols=2, figsize=(6, 3))\n",
        "for iax, (graph_nx, ax) in enumerate(zip(graphs_nx, axs)):\n",
        "  nx.draw(graph_nx, ax=ax)\n",
        "  ax.set_title(\"Graph {}\".format(iax))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o-HixEQ7sU6q"
      },
      "source": [
        "### Print the `GraphsTuple` fields\n",
        "\n",
        "You can also print out the data contained in a `GraphsTuple` by running the cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "akIs5O2b9KOI"
      },
      "outputs": [],
      "source": [
        "def print_graphs_tuple(graphs_tuple):\n",
        "  print(\"Shapes of GraphsTuple's fields:\")\n",
        "  print(graphs_tuple.map(lambda x: x if x is None else x.shape, fields=graphs.ALL_FIELDS))\n",
        "  print(\"\\nData contained in GraphsTuple's fields:\")\n",
        "  print(f\"globals:\\n{graphs_tuple.globals}\")\n",
        "  print(f\"nodes:\\n{graphs_tuple.nodes}\")\n",
        "  print(f\"edges:\\n{graphs_tuple.edges}\")\n",
        "  print(f\"senders:\\n{graphs_tuple.senders}\")\n",
        "  print(f\"receivers:\\n{graphs_tuple.receivers}\")\n",
        "  print(f\"n_node:\\n{graphs_tuple.n_node}\")\n",
        "  print(f\"n_edge:\\n{graphs_tuple.n_edge}\")\n",
        "\n",
        "print_graphs_tuple(graphs_tuple)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3TaUiWVh_SX9"
      },
      "source": [
        "### Back to data dicts\n",
        "It is also possible to retrieve a `list` of  graph `dict`s by using `utils_np.graphs_tuple_to_data_dicts`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PRrsWIDK_Sm5"
      },
      "outputs": [],
      "source": [
        "recovered_data_dict_list = utils_np.graphs_tuple_to_data_dicts(graphs_tuple)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JtUYYeERq2x2"
      },
      "source": [
        "## Ways to represent different data sources with a graph\n",
        "\n",
        "As mentioned above, some graph properties are optional, and the following two cells show two instances of how this can be used."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nz7RDJL1r8Mh"
      },
      "source": [
        "### Graph with no features"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PlBpGYcx0Ux5"
      },
      "outputs": [],
      "source": [
        "# Number of nodes\n",
        "n_node = 3\n",
        "\n",
        "# Three edges connecting the nodes in a cycle\n",
        "senders = [0, 1, 2]  # Indices of nodes sending the edges\n",
        "receivers = [1, 2, 0]  # Indices of nodes receiving the edges\n",
        "\n",
        "data_dict = {\n",
        "    \"n_node\": n_node,\n",
        "    \"senders\": senders,\n",
        "    \"receivers\": receivers,\n",
        "}\n",
        "graphs_tuple = utils_np.data_dicts_to_graphs_tuple([data_dict])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "efQtkXoFr5su"
      },
      "source": [
        "### Set (ie. graph without edges)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J18qD3fU41OE"
      },
      "outputs": [],
      "source": [
        "# Node features.\n",
        "nodes = [[10.],  # Node 0\n",
        "         [11.],  # Node 1\n",
        "         [12.]]  # Node 2\n",
        "\n",
        "data_dict = {\n",
        "    \"nodes\": nodes,\n",
        "}\n",
        "\n",
        "graphs_tuple = utils_np.data_dicts_to_graphs_tuple([data_dict])\n",
        "\n",
        "# We can visualize the graph using networkx.\n",
        "graphs_nx = utils_np.graphs_tuple_to_networkxs(graphs_tuple)\n",
        "ax = plt.figure(figsize=(3, 3)).gca()\n",
        "nx.draw(graphs_nx[0], ax=ax)\n",
        "_ = ax.set_title(\"Graph without edges\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oPrZFvcRr1yU"
      },
      "source": [
        "### Creating a `GraphsTuple` from a `networkx` graph\n",
        "\n",
        "[`networkx`](https://networkx.github.io/) is a powerful graph manipulation library in  Python. A `GraphsTuple` to be built from `networkx` graphs as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ONEOf83YyTKp"
      },
      "outputs": [],
      "source": [
        "graph_nx = nx.OrderedMultiDiGraph()\n",
        "\n",
        "# Globals.\n",
        "graph_nx.graph[\"features\"] = np.array([0.6, 0.7, 0.8])\n",
        "\n",
        "# Nodes.\n",
        "graph_nx.add_node(0, features=np.array([0.3, 1.3]))\n",
        "graph_nx.add_node(1, features=np.array([0.4, 1.4]))\n",
        "graph_nx.add_node(2, features=np.array([0.5, 1.5]))\n",
        "graph_nx.add_node(3, features=np.array([0.6, 1.6]))\n",
        "\n",
        "# Edges.\n",
        "graph_nx.add_edge(0, 1, features=np.array([3.6, 3.7]))\n",
        "graph_nx.add_edge(2, 0, features=np.array([5.6, 5.7]))\n",
        "graph_nx.add_edge(3, 0, features=np.array([6.6, 6.7]))\n",
        "\n",
        "ax = plt.figure(figsize=(3, 3)).gca()\n",
        "nx.draw(graph_nx, ax=ax)\n",
        "ax.set_title(\"Graph\")\n",
        "\n",
        "graphs_tuple = utils_np.networkxs_to_graphs_tuple([graph_nx])\n",
        "\n",
        "print_graphs_tuple(graphs_tuple)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cSUY6pbYXE1i"
      },
      "source": [
        "## Working with tensor `GraphsTuple`'s"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J7tlhBX1Z8RP"
      },
      "outputs": [],
      "source": [
        "#@title Define functions for generating and plotting graphs\n",
        "\n",
        "GLOBAL_SIZE = 4\n",
        "NODE_SIZE = 5\n",
        "EDGE_SIZE = 6\n",
        "\n",
        "def get_graph_data_dict(num_nodes, num_edges):\n",
        "  return {\n",
        "      \"globals\": np.random.rand(GLOBAL_SIZE).astype(np.float32),\n",
        "      \"nodes\": np.random.rand(num_nodes, NODE_SIZE).astype(np.float32),\n",
        "      \"edges\": np.random.rand(num_edges, EDGE_SIZE).astype(np.float32),\n",
        "      \"senders\": np.random.randint(num_nodes, size=num_edges, dtype=np.int32),\n",
        "      \"receivers\": np.random.randint(num_nodes, size=num_edges, dtype=np.int32),\n",
        "  }\n",
        "\n",
        "graph_3_nodes_4_edges = get_graph_data_dict(num_nodes=3, num_edges=4)\n",
        "graph_5_nodes_8_edges = get_graph_data_dict(num_nodes=5, num_edges=8)\n",
        "graph_7_nodes_13_edges = get_graph_data_dict(num_nodes=7, num_edges=13)\n",
        "graph_9_nodes_25_edges = get_graph_data_dict(num_nodes=9, num_edges=25)\n",
        "\n",
        "graph_dicts = [graph_3_nodes_4_edges, graph_5_nodes_8_edges,\n",
        "               graph_7_nodes_13_edges, graph_9_nodes_25_edges]\n",
        "\n",
        "def plot_graphs_tuple(graphs_tuple):\n",
        "  networkx_graphs = utils_np.graphs_tuple_to_networkxs(graphs_tuple)\n",
        "  num_graphs = len(networkx_graphs)\n",
        "  _, axes = plt.subplots(1, num_graphs, figsize=(5*num_graphs, 5))\n",
        "  if num_graphs == 1:\n",
        "    axes = axes,\n",
        "  for graph, ax in zip(networkx_graphs, axes):\n",
        "    plot_graph_networkx(graph, ax)\n",
        "\n",
        "\n",
        "def plot_graph_networkx(graph, ax, pos=None):\n",
        "  node_labels = {node: \"{:.3g}\".format(data[\"features\"][0])\n",
        "                 for node, data in graph.nodes(data=True)\n",
        "                 if data[\"features\"] is not None}\n",
        "  edge_labels = {(sender, receiver): \"{:.3g}\".format(data[\"features\"][0])\n",
        "                 for sender, receiver, data in graph.edges(data=True)\n",
        "                 if data[\"features\"] is not None}\n",
        "  global_label = (\"{:.3g}\".format(graph.graph[\"features\"][0])\n",
        "                  if graph.graph[\"features\"] is not None else None)\n",
        "\n",
        "  if pos is None:\n",
        "    pos = nx.spring_layout(graph)\n",
        "  nx.draw_networkx(graph, pos, ax=ax, labels=node_labels)\n",
        "\n",
        "  if edge_labels:\n",
        "    nx.draw_networkx_edge_labels(graph, pos, edge_labels, ax=ax)\n",
        "\n",
        "  if global_label:\n",
        "    plt.text(0.05, 0.95, global_label, transform=ax.transAxes)\n",
        "\n",
        "  ax.yaxis.set_visible(False)\n",
        "  ax.xaxis.set_visible(False)\n",
        "  return pos\n",
        "\n",
        "\n",
        "def plot_compare_graphs(graphs_tuples, labels):\n",
        "  pos = None\n",
        "  num_graphs = len(graphs_tuples)\n",
        "  _, axes = plt.subplots(1, num_graphs, figsize=(5*num_graphs, 5))\n",
        "  if num_graphs == 1:\n",
        "    axes = axes,\n",
        "  pos = None\n",
        "  for name, graphs_tuple, ax in zip(labels, graphs_tuples, axes):\n",
        "    graph = utils_np.graphs_tuple_to_networkxs(graphs_tuple)[0]\n",
        "    pos = plot_graph_networkx(graph, ax, pos=pos)\n",
        "    ax.set_title(name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_kygzbsSXsHJ"
      },
      "source": [
        "## Creating a constant tensor `GraphsTuple` from data dicts\n",
        "\n",
        "Similar to `utils_np.data_dicts_to_graphs_tuple`,  the `utils_tf` module, which manipulates graphs whose attributes are represented as TensorFlow tensors, contains a function named [`utils_tf.data_dicts_to_graphs_tuple`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#utils_tfgraphs_tuple_to_data_dictsgraph), which creates a constant tensor graph from data dicts, containing either numpy arrays of tensors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yf8d8qV2gQV8"
      },
      "outputs": [],
      "source": [
        "graphs_tuple_tf = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "plot_graphs_tuple(graphs_tuple_tf)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ufxDR3vpuG5"
      },
      "outputs": [],
      "source": [
        "# If the GraphsTuple has None's are automatically handled\n",
        "graphs_tuple_tf = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "\n",
        "# Removing the edges from a graph.\n",
        "graph_with_nones = graphs_tuple_tf.replace(\n",
        "    edges=None, senders=None, receivers=None, n_edge=graphs_tuple_tf.n_edge*0)\n",
        "\n",
        "plot_graphs_tuple(graph_with_nones)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UuBFLcpgvgLL"
      },
      "source": [
        "A similar utility is provided to work with `networkx` graphs: [`utils_np.data_dict_to_networkx`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#utils_npdata_dict_to_networkxdata_dict)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "acSavUGPXsZI"
      },
      "source": [
        "## Slicing graphs from within a batch\n",
        "\n",
        "A subset of graphs can be retrieved from a batch represented by `GraphsTuple` as follows."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "63OYmvUxZBWn"
      },
      "outputs": [],
      "source": [
        "# If the GraphsTuple has None's we need to make use of `utils_tf.make_runnable_in_session`.\n",
        "graphs_tuple_tf = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "\n",
        "first_graph_tf = utils_tf.get_graph(graphs_tuple_tf, 0)\n",
        "three_graphs_tf = utils_tf.get_graph(graphs_tuple_tf, slice(1, 4))\n",
        "\n",
        "plot_graphs_tuple(first_graph_tf)\n",
        "plot_graphs_tuple(three_graphs_tf)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MLD2KEdEYldQ"
      },
      "source": [
        "## Concatenating multiple `GraphsTuple` instances"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MrWouGJHxeIO"
      },
      "outputs": [],
      "source": [
        "# Concatenating along the batch dimension\n",
        "graphs_tuple_1_tf = utils_tf.data_dicts_to_graphs_tuple(graph_dicts[0:1])\n",
        "graphs_tuple_2_tf = utils_tf.data_dicts_to_graphs_tuple(graph_dicts[1:])\n",
        "\n",
        "graphs_tuple_tf = utils_tf.concat([graphs_tuple_1_tf, graphs_tuple_2_tf], axis=0)\n",
        "\n",
        "plot_graphs_tuple(graphs_tuple_tf)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "btcMDrixx9Bs"
      },
      "source": [
        "Similarly, we can concatenate along feature dimensions, assuming all of the batches to be concatenates have the same graph structure/connectivity."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_KzQhLJJYEQZ"
      },
      "source": [
        "See `utils_tf` for more methods to work with GraphsTuple's containing tensors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VUzXQDdRaIKt"
      },
      "source": [
        "# Graph Net modules\n",
        "\n",
        "So far we've covered how to manipulate graph-structured data via the `GraphsTuple` class and the `utils_np` and `utils_tf` utilities. Now we show how to use actual graph networks."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0rOU3rjRbK9O"
      },
      "source": [
        "## Creating a [`modules.GraphNetwork`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-modulesgraphnetwork)\n",
        "\n",
        "![Graph Network block](https://github.com/deepmind/graph_nets/raw/master/images/graph-network.png)\n",
        "\n",
        "A graph network has up to three learnable sub-functions: edge ($\\phi^e$), node ($\\phi^v$), and global ($\\phi^u$) in the schematic above. See Section 3.2.2, page 12 in [our companion arXiv paper](https://arxiv.org/pdf/1806.01261.pdf) for details.\n",
        "\n",
        "To instantiate a graph network module in the library, these sub-functions are specified via constructor arguments which are `callable`s that return Sonnet modules, such as `snt.Linear` or `snt.nets.MLP`.\n",
        "\n",
        "The reason that a `callable` is provided, instead of the module/method directly, is so the Graph Net object owns the modules and the variables created by them."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P5_ADeDsyQNZ"
      },
      "outputs": [],
      "source": [
        "OUTPUT_EDGE_SIZE = 10\n",
        "OUTPUT_NODE_SIZE = 11\n",
        "OUTPUT_GLOBAL_SIZE = 12\n",
        "\n",
        "graph_network = modules.GraphNetwork(\n",
        "    edge_model_fn=lambda: snt.Linear(output_size=OUTPUT_EDGE_SIZE),\n",
        "    node_model_fn=lambda: snt.Linear(output_size=OUTPUT_NODE_SIZE),\n",
        "    global_model_fn=lambda: snt.Linear(output_size=OUTPUT_GLOBAL_SIZE))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VsSIrvMEbV1W"
      },
      "source": [
        "## Feeding a `GraphsTuple` to a Graph Net\n",
        "\n",
        "A `GraphsTuple` can be fed into a graph network, which returns an output graph with the same number of nodes, edges, and edge connectivity, but with updated edge, node and global features. All of the output features are conditioned on the input features according to the graph structure, and are fully differentiable."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "10g2R-aKz2kY"
      },
      "outputs": [],
      "source": [
        "input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "output_graphs = graph_network(input_graphs)\n",
        "\n",
        "print(f\"Output edges size: {output_graphs.edges.shape[-1]}\")  # Equal to OUTPUT_EDGE_SIZE\n",
        "print(f\"Output nodes size: {output_graphs.nodes.shape[-1]}\")  # Equal to OUTPUT_NODE_SIZE\n",
        "print(f\"Output globals size: {output_graphs.globals.shape[-1]}\")  # Equal to OUTPUT_GLOBAL_SIZE"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HGK5xGBob3lb"
      },
      "source": [
        "## Connecting a `GraphNetwork` recurrently\n",
        "\n",
        "A Graph Net module can be chained recurrently by matching the output feature sizes to the input feature sizes, and feeding the output back to the input multiple times ([arXiv paper](https://arxiv.org/abs/1806.01261), bottom of Fig. 6a)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e54GDFXKB2NO"
      },
      "outputs": [],
      "source": [
        "input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "\n",
        "graph_network = modules.GraphNetwork(\n",
        "    edge_model_fn=lambda: snt.Linear(output_size=EDGE_SIZE),\n",
        "    node_model_fn=lambda: snt.Linear(output_size=NODE_SIZE),\n",
        "    global_model_fn=lambda: snt.Linear(output_size=GLOBAL_SIZE))\n",
        "\n",
        "num_recurrent_passes = 3\n",
        "previous_graphs = input_graphs\n",
        "for unused_pass in range(num_recurrent_passes):\n",
        "  previous_graphs = graph_network(previous_graphs)\n",
        "output_graphs = previous_graphs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kZzEBReqCimx"
      },
      "source": [
        "Alternatively, we can process the input graph multiple times with a graph state that gets updated recurrently."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_rfSo5NJ1UwG"
      },
      "outputs": [],
      "source": [
        "def zeros_graph(sample_graph, edge_size, node_size, global_size):\n",
        "  zeros_graphs = sample_graph.replace(nodes=None, edges=None, globals=None)\n",
        "  zeros_graphs = utils_tf.set_zero_edge_features(zeros_graphs, edge_size)\n",
        "  zeros_graphs = utils_tf.set_zero_node_features(zeros_graphs, node_size)\n",
        "  zeros_graphs = utils_tf.set_zero_global_features(zeros_graphs, global_size)\n",
        "  return zeros_graphs\n",
        "\n",
        "graph_network = modules.GraphNetwork(\n",
        "    edge_model_fn=lambda: snt.Linear(output_size=OUTPUT_EDGE_SIZE),\n",
        "    node_model_fn=lambda: snt.Linear(output_size=OUTPUT_NODE_SIZE),\n",
        "    global_model_fn=lambda: snt.Linear(output_size=OUTPUT_GLOBAL_SIZE))\n",
        "\n",
        "input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "initial_state = zeros_graph(\n",
        "    input_graphs, OUTPUT_EDGE_SIZE, OUTPUT_NODE_SIZE, OUTPUT_GLOBAL_SIZE)\n",
        "\n",
        "num_recurrent_passes = 3\n",
        "\n",
        "current_state = initial_state\n",
        "for unused_pass in range(num_recurrent_passes):\n",
        "  input_and_state_graphs = utils_tf.concat(\n",
        "      [input_graphs, current_state], axis=1)\n",
        "  current_state = graph_network(input_and_state_graphs)\n",
        "output_graphs = current_state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PNnLXPXTmZXW"
      },
      "source": [
        "Similarly, recurrent modules with gating, such as an LSTM or GRU, can be applied on the edges, nodes, and globals of the state and input graphs separately."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6m6aOhOtAz3T"
      },
      "source": [
        "## Other canonical Graph Net modules\n",
        "\n",
        "Other canonical modules discussed in Figure 4 of our [arXiv paper](https://arxiv.org/abs/1806.01261) are provided in `graph_nets.modules`:\n",
        "- [`modules.GraphIndependent`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-modulesgraphindependent) (updates the global, node, and edge features independently, without message-passing)\n",
        "- [`modules.InteractionNetwork`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-modulesinteractionnetwork) (an example of a \"Message-passing neural network\")\n",
        "- [`modules.CommNet`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-modulescommnet) (another example of a \"Message-passing neural network\")\n",
        "- [`modules.SelfAttention`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-modulesselfattention) (an example of a \"Non-local neural network\")\n",
        "- [`modules.RelationNetwork`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-modulesrelationnetwork)\n",
        "- [`modules.DeepSets`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-modulesdeepsets)\n",
        "\n",
        "See documentation for more details and corresponding references."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qfhlV2qpah5O"
      },
      "source": [
        "# Graph Net building blocks\n",
        "\n",
        "Custom graph net modules can be built using a few basic building blocks provided in `graph_nets.blocks`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ta0YenqKaq7E"
      },
      "source": [
        "## Broadcast operations\n",
        "\n",
        "Broadcast operations allow to transfer information between different types of elements in the graph:\n",
        "- [`blocks.broadcast_globals_to_nodes`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#blocksbroadcast_globals_to_nodesgraph-namebroadcast_globals_to_nodes): Copy/broadcast global features across all nodes.\n",
        "- [`blocks.broadcast_globals_to_edges`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#blocksbroadcast_globals_to_edgesgraph-namebroadcast_globals_to_edges): Copy/broadcast global features across all edges.\n",
        "- [`blocks.broadcast_sender_nodes_to_edges`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#blocksbroadcast_sender_nodes_to_edgesgraph-namebroadcast_sender_nodes_to_edges): Copy/broadcast node information from each node, across all edges for which that node is a sender.\n",
        "- [`blocks.broadcast_receiver_nodes_to_edges`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#blocksbroadcast_receiver_nodes_to_edgesgraph-namebroadcast_receiver_nodes_to_edges): Copy/broadcast node information from each node, across all edges for which that node is a receiver."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sR6ckNvQbgbO"
      },
      "outputs": [],
      "source": [
        "graphs_tuple = utils_tf.data_dicts_to_graphs_tuple([data_dict_0])\n",
        "updated_broadcast_globals_to_nodes = graphs_tuple.replace(\n",
        "    nodes=blocks.broadcast_globals_to_nodes(graphs_tuple))\n",
        "updated_broadcast_globals_to_edges = graphs_tuple.replace(\n",
        "    edges=blocks.broadcast_globals_to_edges(graphs_tuple))\n",
        "updated_broadcast_sender_nodes_to_edges = graphs_tuple.replace(\n",
        "    edges=blocks.broadcast_sender_nodes_to_edges(graphs_tuple))\n",
        "updated_broadcast_receiver_nodes_to_edges = graphs_tuple.replace(\n",
        "    edges=blocks.broadcast_receiver_nodes_to_edges(graphs_tuple))\n",
        "\n",
        "output_graphs = [\n",
        "      graphs_tuple,\n",
        "      updated_broadcast_globals_to_nodes,\n",
        "      updated_broadcast_globals_to_edges,\n",
        "      updated_broadcast_sender_nodes_to_edges,\n",
        "      updated_broadcast_receiver_nodes_to_edges]\n",
        "\n",
        "plot_compare_graphs(output_graphs, labels=[\n",
        "    \"Input graph\",\n",
        "    \"blocks.broadcast_globals_to_nodes\",\n",
        "    \"blocks.broadcast_globals_to_edges\",\n",
        "    \"blocks.broadcast_sender_nodes_to_edges\",\n",
        "    \"blocks.broadcast_receiver_nodes_to_edges\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ykPbYewdLdKM"
      },
      "source": [
        "We can easily use broadcasters to, for example, set the value of each edge to be the sum of the first feature element of: the input edges, the sender nodes, the receiver nodes, and the global feature.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w3eWOCl18Q5R"
      },
      "outputs": [],
      "source": [
        "graphs_tuple = utils_tf.data_dicts_to_graphs_tuple([data_dict_0])\n",
        "\n",
        "updated_graphs_tuple = graphs_tuple.replace(\n",
        "    edges=(graphs_tuple.edges[:, :1] +\n",
        "           blocks.broadcast_receiver_nodes_to_edges(graphs_tuple)[:, :1] +\n",
        "           blocks.broadcast_sender_nodes_to_edges(graphs_tuple)[:, :1] +\n",
        "           blocks.broadcast_globals_to_edges(graphs_tuple)[:, :1]))\n",
        "\n",
        "output_graphs = [\n",
        "    graphs_tuple,\n",
        "    updated_graphs_tuple]\n",
        "\n",
        "plot_compare_graphs(output_graphs, labels=[\n",
        "    \"Input graph\",\n",
        "    \"Updated graph\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uE8bBiLbaoCX"
      },
      "source": [
        "## Aggregators\n",
        "\n",
        "Aggregators perform reduce operations between different elements of the graph:\n",
        "\n",
        "- [`blocks.EdgesToGlobalsAggregator`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-blocksedgestoglobalsaggregator): Aggregates the sets of features for all edges into a single global set of features.\n",
        "- [`blocks.NodesToGlobalsAggregator`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-blocksnodestoglobalsaggregator): Aggregates the sets of features for all nodes into a single global set of features.\n",
        "- [`blocks.SentEdgesToNodesAggregator`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-blockssentedgestonodesaggregator): Aggregates the sets of features for all edges sent by each node into a single set of features for that node.\n",
        "- [`blocks.ReceivedEdgesToNodesAggregator`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-blocksreceivededgestonodesaggregator): Aggregates the sets of features for all edges received by each node into a single set of features for that.\n",
        "\n",
        "Different types of reduce operations are:\n",
        "- [`tf.unsorted_segment_sum`](https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_sum): Elementwise sum. Set to 0 for empty sets.\n",
        "- [`tf.unsorted_segment_mean`](https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_mean): Elementwise mean. Set to 0 for empty sets.\n",
        "- [`tf.unsorted_segment_prod`](https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_prod): Elementwise prod. Set to 1 for empty sets.\n",
        "- [`blocks.unsorted_segment_max_or_zero`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#blocksunsorted_segment_max_or_zerovalues-indices-num_groups-nameunsorted_segment_max_or_zero): Elementwise max. Set to 0 for empty sets.\n",
        "- [`blocks.unsorted_segment_min_or_zero`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#blocksunsorted_segment_max_or_zerovalues-indices-num_groups-nameunsorted_segment_min_or_zero): Elementwise min. Set to 0 for empty sets.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qR_3QHNgiMYv"
      },
      "outputs": [],
      "source": [
        "graphs_tuple = utils_tf.data_dicts_to_graphs_tuple([data_dict_0])\n",
        "\n",
        "reducer = tf.math.unsorted_segment_sum\n",
        "updated_edges_to_globals = graphs_tuple.replace(\n",
        "    globals=blocks.EdgesToGlobalsAggregator(reducer=reducer)(graphs_tuple))\n",
        "updated_nodes_to_globals = graphs_tuple.replace(\n",
        "    globals=blocks.NodesToGlobalsAggregator(reducer=reducer)(graphs_tuple))\n",
        "updated_sent_edges_to_nodes = graphs_tuple.replace(\n",
        "    nodes=blocks.SentEdgesToNodesAggregator(reducer=reducer)(graphs_tuple))\n",
        "updated_received_edges_to_nodes = graphs_tuple.replace(\n",
        "    nodes=blocks.ReceivedEdgesToNodesAggregator(reducer=reducer)(graphs_tuple))\n",
        "\n",
        "output_graphs = [\n",
        "    graphs_tuple,\n",
        "    updated_edges_to_globals,\n",
        "    updated_nodes_to_globals,\n",
        "    updated_sent_edges_to_nodes,\n",
        "    updated_received_edges_to_nodes]\n",
        "\n",
        "plot_compare_graphs(output_graphs, labels=[\n",
        "    \"Input graph\",\n",
        "    \"blocks.EdgesToGlobalsAggregator\",\n",
        "    \"blocks.NodesToGlobalsAggregator\",\n",
        "    \"blocks.SentEdgesToNodesAggregator\",\n",
        "    \"blocks.ReceivedEdgesToNodesAggregator\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lE0e3bMna96z"
      },
      "source": [
        "## [`blocks.EdgeBlock`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-blocksedgeblock)\n",
        "\n",
        "An EdgeBlock consists of applying a function to the concatenation of:\n",
        "-  `graphs_tuple.edges`\n",
        "- `blocks.broadcast_sender_nodes_to_edges(graphs_tuple)`\n",
        "- `blocks.broadcast_receiver_nodes_to_edges(graphs_tuple)`\n",
        "- `blocks.broadcast_globals_to_edges(graphs_tuple)`\n",
        "\n",
        "The result is a graph with new edge features conditioned on input edges, nodes and global features according to the graph structure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b549fO5bnxNy"
      },
      "outputs": [],
      "source": [
        "edge_block = blocks.EdgeBlock(\n",
        "    edge_model_fn=lambda: snt.Linear(output_size=10))\n",
        "\n",
        "input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "output_graphs = edge_block(input_graphs)\n",
        "\n",
        "print(f\"Output edges size: {output_graphs.edges.shape[-1]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FsqdEWjLa0kF"
      },
      "source": [
        "## [`blocks.NodeBlock`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-blocksnodeblock)\n",
        "\n",
        "An NodeBlock consists of applying a function to the concatenation of:\n",
        "-  `graphs_tuple.nodes`\n",
        "- `blocks.ReceivedEdgesToNodesAggregator(\u003creducer-function\u003e)(graphs_tuple)`\n",
        "- `blocks.broadcast_globals_to_nodes(graphs_tuple)`\n",
        "\n",
        "The result is a graph with new node features conditioned on input edges, nodes and global features according to the graph structure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LvsrYPqgpZ5x"
      },
      "outputs": [],
      "source": [
        "node_block = blocks.NodeBlock(\n",
        "    node_model_fn=lambda: snt.Linear(output_size=15))\n",
        "\n",
        "input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "output_graphs = node_block(input_graphs)\n",
        "\n",
        "print(f\"Output nodes size: {output_graphs.nodes.shape[-1]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JlNdE_rbbBOj"
      },
      "source": [
        "## [`blocks.GlobalBlock`](https://github.com/deepmind/graph_nets/blob/master/docs/graph_nets.md#class-blocksglobalblock)\n",
        "\n",
        "An GlobalBlock consists of applying a function to the concatenation of:\n",
        "\n",
        "- `graphs_tuple.globals`\n",
        "- `blocks.EdgesToGlobalsAggregator(\u003creducer-function\u003e)(graphs_tuple)`\n",
        "- `blocks.NodesToGlobalsAggregator(\u003creducer-function\u003e)(graphs_tuple)`\n",
        "\n",
        "The result is a graph with new globals features conditioned on input edges, nodes and global features."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sVuZ4EUzph1a"
      },
      "outputs": [],
      "source": [
        "global_block = blocks.GlobalBlock(\n",
        "    global_model_fn=lambda: snt.Linear(output_size=20))\n",
        "\n",
        "input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "output_graphs = global_block(input_graphs)\n",
        "\n",
        "print(f\"Output globals size: {output_graphs.globals.shape[-1]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZC_Y6dEZqNCx"
      },
      "source": [
        "## Block compositionality\n",
        "\n",
        "A `modules.GraphNetwork` is composed internally of a `modules.EdgeBlock`,  a `modules.NodeBlock`, and a `modules.GlobalBlock`, connected serially. This allows every field in the output to be conditioned by any field in the input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eN4Nz9bspqaQ"
      },
      "outputs": [],
      "source": [
        "graph_network = modules.GraphNetwork(\n",
        "    edge_model_fn=lambda: snt.Linear(output_size=10),\n",
        "    node_model_fn=lambda: snt.Linear(output_size=15),\n",
        "    global_model_fn=lambda: snt.Linear(output_size=20))\n",
        "\n",
        "input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)\n",
        "output_graphs = graph_network(input_graphs)\n",
        "\n",
        "for var in graph_network.variables:\n",
        "  print(var)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vGKTIwF6qtfR"
      },
      "source": [
        "Most of the existing neural networks operating on graphs can be built upon this set of building blocks using their different configuration options. See `graph_nets.modules` for some examples."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "graph_nets_basics.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    },
    "toc": {
      "base_numbering": 1,
      "nav_menu": {},
      "number_sections": true,
      "sideBar": true,
      "skip_h1_title": false,
      "title_cell": "Table of Contents",
      "title_sidebar": "Contents",
      "toc_cell": false,
      "toc_position": {},
      "toc_section_display": true,
      "toc_window_display": false
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
